`include "common.svh"

module fetch #(
    parameter FETCH_WIDTH = 2,
    localparam BPU_AW = $clog2(FETCH_WIDTH)
) (
    input clk,
    input rst,
    // Redirect Req by backend
    input i_redirect_valid,
    input word_t i_redirect_target,
    // IRAM Req
    input o_iram_ready,
    output o_iram_valid,
    output IRAM_Req o_iram_req,
    // IRAM Rsp
    input IRAM_Rsp i_iram_rsp,
    // BPU Req
    output o_bpu_valid,
    input o_bpu_ready,
    output word_t o_bpu_pc,
    // BPU Rsp
    input i_bpu_valid,
    input i_bpu_res_valid,
    input BPU_Res i_bpu_res,
    input [BPU_AW-1:0] i_bpu_res_offset,
    input i_bpu_res_multi_branch,
    // IBUF
    output o_ibuf_valid,
    input o_ibuf_ready,
    output o_ibuf_entry_valid[FETCH_WIDTH-1:0],
    output IBUF_Entry o_ibuf_entry[FETCH_WIDTH-1:0]
);

  wire o_bpu_fire = o_bpu_valid & o_bpu_ready;
  wire o_iram_fire = o_iram_valid & o_iram_ready;
  wire o_ibuf_fire = o_ibuf_valid & o_ibuf_ready;

  logic s1_ready;

  wire flush = i_redirect_valid;
  word_t pc, next_pc;

  word_t s1_bpu_target;
  word_t s1_bpu_target_next;
  logic s1_bpu_taken;
  // Stage 0 Caculate fetch PC to Request BPU and iram
  wire s0_fire = o_bpu_fire & o_iram_fire;
  wire s0_valid = (~flush) & s1_ready;
  word_t s0_fetch_pc, s0_fetch_next_pc;
  assign s0_fetch_pc = s1_bpu_taken ? s1_bpu_target : pc;
  assign s0_fetch_next_pc = s1_bpu_taken ? s1_bpu_target_next : pc + (FETCH_WIDTH * 4);
  assign next_pc = i_redirect_valid ? i_redirect_target : s0_fetch_next_pc;

  wire pc_load = (s1_ready & s0_fire) | i_redirect_valid;
  reg_lr #(`WORD_BITS, `RESET_PC) pc_r (
      .load(pc_load),
      .din (next_pc),
      .dout(pc),
      .*
  );

  assign o_iram_valid = s0_valid & o_bpu_ready & s1_ready;
  assign o_iram_req.pc = s0_fetch_pc;
  // ensure that a pc only requests BPU once
  assign o_bpu_valid = s0_valid & o_iram_ready & s1_ready;
  assign o_bpu_pc = s0_fetch_pc;

  // Stage 1 Caculate fetch PC and next PC and latch iram_data and bpu_res
  assign s1_ready = s1_valid ? o_ibuf_ready : 1'b1;
  logic s1_valid;
  wire s1_valid_set = s0_fire & s1_ready;
  wire s1_valid_clr = o_ibuf_fire;
  reg_lr #(1) s1_valid_r (
      .rst (rst | flush),
      .load(s1_valid_set | s1_valid_clr),
      .din (s1_valid_set),
      .dout(s1_valid),
      .*
  );
  word_t s1_fetch_pc;
  reg_l #(`WORD_BITS) s1_fetch_pc_r (
      .load(s1_ready & s0_fire),
      .din (s0_fetch_pc),
      .dout(s1_fetch_pc),
      .*
  );

  // latch bpu res
  wire latch_bpu = ~i_bpu_valid;
  logic s1_bpu_res_valid, i_bpu_res_valid_r;
  logic s1_bpu_res_multi_branch, i_bpu_res_multi_branch_r;
  BPU_Res s1_bpu_res, i_bpu_res_r;
  logic [BPU_AW-1:0] s1_bpu_res_offset, i_bpu_res_offset_r;

  assign s1_bpu_res_valid = latch_bpu ? i_bpu_res_valid_r : i_bpu_res_valid;
  assign s1_bpu_res = latch_bpu ? i_bpu_res_r : i_bpu_res;
  assign s1_bpu_res_offset = latch_bpu ? i_bpu_res_offset_r : i_bpu_res_offset;
  assign s1_bpu_res_multi_branch = latch_bpu ? i_bpu_res_multi_branch_r : i_bpu_res_multi_branch;

  reg_lr #($size(
      BPU_Res
  ) + BPU_AW + 2) bpu_res_r (
      .rst (rst | flush),
      .load(~latch_bpu),
      .din ({i_bpu_res_valid, i_bpu_res, i_bpu_res_offset, i_bpu_res_multi_branch}),
      .dout({i_bpu_res_valid_r, i_bpu_res_r, i_bpu_res_offset_r, i_bpu_res_multi_branch_r}),
      .*
  );

  assign s1_bpu_taken = (s1_bpu_res_multi_branch | s1_bpu_res.taken) & s1_bpu_res_valid;
  assign s1_bpu_target = s1_bpu_res.target;
  assign s1_bpu_target_next = s1_bpu_target + (FETCH_WIDTH * 4);

  // latch iram data
  IRAM_Rsp iram_rsp, s1_iram_rsp;
  wire latch_iram = ~o_iram_ready;
  assign s1_iram_rsp = latch_iram ? iram_rsp : i_iram_rsp;
  reg_lr #($size(
      IRAM_Rsp
  )) iram_rsp_r (
      .rst (rst | flush),
      .load(~latch_iram),
      .din (i_iram_rsp),
      .dout(iram_rsp),
      .*
  );
  /* verilator lint_off UNSIGNED */
  genvar i;
  generate
    for (i = 0; i < FETCH_WIDTH; i = i + 1) begin
      logic [BPU_AW-1:0] index = i;  // i < offset ? 1 : taken
      assign o_ibuf_entry_valid[i] = (index <= s1_bpu_res_offset) ? 1'b1 : (~s1_bpu_taken);
      assign o_ibuf_entry[i].pc = s1_fetch_pc + (i * 4);
      assign o_ibuf_entry[i].inst = s1_iram_rsp.inst[i];
      assign o_ibuf_entry[i].bpu_res = s1_bpu_res;
      assign o_ibuf_entry[i].bpu_res_valid = s1_bpu_res_valid & (s1_bpu_res_offset == index);
    end
  endgenerate
  assign o_ibuf_valid = s1_valid;
  /* verilator lint_on UNSIGNED */
endmodule
